import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split, cross_val_predict
from sklearn.linear_model import LogisticRegression
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics import classification_report
from sklearn.pipeline import Pipeline
from sklearn.model_selection import GridSearchCV

# 加载数据
train_data = pd.read_csv('ChatGPT生成文本检测器公开数据-更新/train.csv')
test_data = pd.read_csv('ChatGPT生成文本检测器公开数据-更新/test.csv')

# 数据预处理
train_data['content'] = train_data['content'].apply(lambda x: x[1:-1])
test_data['content'] = test_data['content'].apply(lambda x: x[1:-1])

# 划分训练集和验证集
train_text, valid_text, train_label, valid_label = train_test_split(
    train_data['content'], train_data['label'], test_size=0.2, random_state=42
)

# 使用Pipeline进行流水线构建
pipeline = Pipeline([
    ('tfidf', TfidfVectorizer(token_pattern=r'\w{1,}', max_features=5000, ngram_range=(1, 2))),
    ('model', LogisticRegression(max_iter=1000))
])

# 网格搜索参数
param_grid = {
    'tfidf__max_features': [2000, 5000],
    'model__C': [0.1, 1, 10]
}

# 在训练集上进行网格搜索交叉验证
grid_search = GridSearchCV(pipeline, param_grid, cv=3, scoring='f1_macro')
grid_search.fit(train_text, train_label)

# 最佳模型
best_model = grid_search.best_estimator_

# 打印最佳参数
print("Best Parameters:", grid_search.best_params_)

# 预测验证集
valid_predictions = best_model.predict(valid_text)

# 打印验证集的分类报告
print("Validation Set Classification Report:")
print(classification_report(valid_label, valid_predictions))

# 使用最佳模型进行测试集预测
test_predictions = best_model.predict(test_data['content'])
test_data['label'] = test_predictions

# 保存预测结果
test_data[['name', 'label']].to_csv('tfidf_predictions.csv', index=None)
